# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: xguo

import datetime
import hashlib
import hmac
import json
import logging
import time
import zlib
import urllib.parse
import base64

from tornado import gen
from tornado.ioloop import IOLoop
from tornado.websocket import websocket_connect, WebSocketClosedError
from tornado.concurrent import Future

from coin.base.timestamp import get_timestamp
from coin.exchange.base.rest.task_scheduler import Scheduler
from coin.exchange.base.kr_rest.auth_base import AuthKey
from coin.exchange.huobi.ws.ws_private_client import OperatorType

HUOBI_FUTURES_WS_API_HOST = 'api.hbdm.com'


class HuobiFuturesWsPrivateClient(object):
  def __init__(
      self,
      key_file,
      *,
      ioloop=None,
      notification_callback=None,
      close_callback=None,
      on_ready=None,
      timeout_sec=10,
      api_host=None,
      disable_ssl=False,
  ):
    self._ioloop = ioloop or IOLoop.current()
    scheme = 'ws://' if disable_ssl else 'wss://'
    api_host = api_host or HUOBI_FUTURES_WS_API_HOST
    self._api_url = scheme + api_host + '/notification'
    self._ws = None
    self._ready = False
    self._timeout = timeout_sec * 1e9
    self._num_timeouts = 0
    self._key = AuthKey.from_file(key_file)

    self._notification_callback = notification_callback or (lambda x: None)
    self._close_callback = close_callback or (lambda: None)
    self._on_ready = on_ready or (lambda: None)
    self._scheduler = Scheduler()

    self._last_received_message_timestamp = None
    self._pending_future = None
    self._subscribed_topics = set()
    self._request_id = int(time.time() * 1000)

    self._stopped = False
    self._init_params = (ioloop, notification_callback, close_callback, timeout_sec)
    self._logger = logging.getLogger(__name__)

  @property
  def last_received_message_timestamp(self):
    return self._last_received_message_timestamp

  def is_ready(self):
    return self._ready

  def _get_next_request_id(self):
    self._request_id += 1
    return str(self._request_id)

  def _check_future_timeout(self):
    timestamp = get_timestamp()
    if self._pending_future is None:
      return
    if timestamp - self._pending_future.timestamp > self._timeout:
      self._num_timeouts += 1
      self._pending_future.set_exception(TimeoutError('Websocket request timeout!'))

  def _check_connection(self):
    if self._last_received_message_timestamp is None:
      return

    timestamp = get_timestamp()
    diff_time = timestamp - self._last_received_message_timestamp
    if diff_time > self._timeout:
      self._logger.warning('No message received for %s seconds', diff_time / 1e9)
      self.stop()

  def _gen_auth_msg(self):
    params = {
        'AccessKeyId': self._key.access_key,
        'SignatureMethod': 'HmacSHA256',
        'SignatureVersion': '2',
        'Timestamp': datetime.datetime.utcnow().strftime('%Y-%m-%dT%H:%M:%S'),
    }
    sorted_params = sorted(params.items(), key=lambda x: x[0])
    query_string = urllib.parse.urlencode(sorted_params)
    parsed_results = urllib.parse.urlparse(self._api_url)
    if parsed_results.netloc == 'xiguaapi.hbdm.com':
      api_host = 'api.hbdm.com'
    else:
      api_host = parsed_results.netloc
    payload = ['GET', api_host, parsed_results.path, query_string]
    signature = hmac.new(self._key.secret_key_bytes,
                         msg='\n'.join(payload).encode(),
                         digestmod=hashlib.sha256).digest()
    params['op'] = OperatorType.AUTH.value
    params['type'] = 'api'
    params['Signature'] = base64.b64encode(signature).decode()
    return params

  @gen.coroutine
  def prepare_subscription(self, symbols):
    for symbol in symbols:
      response = yield self.subscribe_orders(symbol)
      if response['err-code'] != 0:
        self.stop()
        return
      else:
        self._logger.info('Subscribe to huobi order %s, %s', symbol, response)

  @gen.coroutine
  def _connect(self):
    # noinspection PyBroadException
    try:
      self._logger.info('%s try to connect to the websocket (%s).',
                        type(self).__name__,
                        self._api_url)
      self._ws = yield websocket_connect(self._api_url)
      self._ioloop.add_callback(self._run)

      self._logger.info('Send auth info!')
      response = yield self.send_request(operation=OperatorType.AUTH)
      if response['err-code'] != 0:
        self._logger.error('Authentication failed! %s', response)
        self.stop()
        return
      else:
        self._logger.info('Huobi Futures WS API Successfully authenticated!')

      self._ready = True
      self._on_ready()
      self._logger.info('Huobi Futures private client websocket connected.')
    except Exception:
      self._logger.exception('%s connect to websocket failed!', type(self).__name__)
      self.stop()

  @gen.coroutine
  def _run(self):
    while not self._stopped:
      msg = None
      # noinspection PyBroadException
      try:
        msg = yield self._ws.read_message()
      except Exception:
        self._logger.exception('Error reading ws message')
        self.stop()

      if msg is None:
        self._logger.error('Websocket is closed!')
        self.stop()
        return

      if len(msg) == 0:
        self._logger.info('Empty message received!')
        continue

      # noinspection PyBroadException
      try:
        self.on_message(msg)
      except Exception:
        self._logger.exception('Exception raised when process message!')

  def on_message(self, raw_msg: bytes):
    timestamp = get_timestamp()
    self._last_received_message_timestamp = timestamp
    if self._stopped:
      return

    try:
      data = json.loads(zlib.decompress(raw_msg, wbits=31))
    except ValueError:
      self._logger.exception('Invalid json string!, %s', raw_msg)
      return

    self._logger.debug('Receive ws msg: %s\n', data)
    operation = data.get('op')
    if operation == OperatorType.PING.value:
      self.send_request(OperatorType.PONG, ts=data['ts'])
    elif operation == OperatorType.PONG.value:
      self._logger.debug('pong received! %s', data)
    elif operation in (OperatorType.AUTH.value,
                       OperatorType.REQ.value,
                       OperatorType.SUB.value,
                       OperatorType.UNSUB.value):

      topic = data.get('topic')
      if operation == OperatorType.SUB and topic:
        self._subscribed_topics.add(topic)

      fut = self._pending_future
      self._pending_future = None
      if fut is None:
        self._logger.error('Unexpected message! %s', data)
      else:
        fut.set_result(data)
    elif operation == OperatorType.NOTIFY.value:
      self._notification_callback(data)
    elif operation == OperatorType.CLOSE.value:
      self.stop()
    elif operation == OperatorType.ERROR.value:
      self._logger.error('Error message received!\n%s', data)
    else:  # message from subscription stream
      self._logger.warning('Not supported yet!', data)

  def _write_message(self, msg):
    try:
      self._ws.write_message(json.dumps(msg))
    except WebSocketClosedError as e:
      self._logger.error('Write message error, ws is close!')
      if self._pending_future is not None:
        self._pending_future.set_exception(e)
        self._pending_future = None

  def send_request(self, operation, *, topic=None, ts=None):
    if self._ws is None:
      self._logger.error('websocket not connected yet! operation=%s, params=%s',
                         operation.value,
                         topic)
      return

    fut = None
    request_id = None
    if operation not in (OperatorType.PING, OperatorType.PONG):
      fut = Future()
      fut.timestamp = get_timestamp()
      fut.operation = operation
      request_id = self._get_next_request_id()
      self._pending_future = fut

    if operation == OperatorType.AUTH:
      request_msg = self._gen_auth_msg()
      request_msg['cid'] = request_id
    elif operation == OperatorType.PONG:
      request_msg = {
          'op': operation.value,
          'ts': ts,
      }
    elif operation == OperatorType.PING:
      request_msg = {
          'op': operation.value,
          'ts': int(time.time() * 1000),
      }
    else:
      assert topic is not None
      request_msg = {
          'op': operation.value,
          'topic': topic,
          'cid': request_id,
      }

    self._logger.debug('send request: %s\n', request_msg)
    self._ioloop.add_callback(self._write_message, request_msg)
    return fut

  def subscribe_topic(self, topic):
    operation = OperatorType.SUB
    fut = self.send_request(operation, topic=topic)
    return fut

  def unsubscribe_topic(self, topic):
    operation = OperatorType.UNSUB
    fut = self.send_request(operation, topic=topic)
    return fut

  def subscribe_orders(self, symbol):
    topic = 'orders.{symbol}'.format(symbol=symbol)
    fut = self.subscribe_topic(topic)
    return fut

  def unsubscribe_orders(self, symbol):
    topic = 'orders.{symbol}'.format(symbol=symbol)
    fut = self.unsubscribe_topic(topic)
    return fut

  def start(self):
    self._ioloop.add_callback(self._connect)
    self._scheduler.add_periodic_task(1, self._check_future_timeout)
    self._scheduler.add_periodic_task(10, self._check_connection)
    self._scheduler.start()

  def stop(self):
    if self._stopped:
      return
    self._scheduler.stop()
    self._ready = False
    self._ioloop.add_callback(self._async_stop)
    self._logger.info('Huobi ws private client (%s) stopped!', id(self))

  @gen.coroutine
  def _async_stop(self):
    for topic in self._subscribed_topics:
      response = yield self.unsubscribe_topic(topic)
      self._logger.info('Topic (%s) unsubscribed!\n%s', topic, response)

    self._stopped = True
    if self._ws is not None:
      self._ws.close()
      self._ws = None
    yield gen.sleep(10)
    self._close_callback()

  def clone(self):
    return type(self)(*self._init_params)
